package tree;

import java.io.*;
import java.util.*;

/**
 * @description: 赫弗曼树
 * 缺陷
 * 本博文所采用的压缩算法是在内存中进行转换后全部存储到一个字符串对象中（未使用缓冲器 一个字节一个字节解压 效率非常低
 * ,如果压缩文件压缩反而更大或者大小没变
 * ，可能文件已经是被压缩的使用哈夫曼压缩算法不起作用），
 * 然后再去写入文件，解压算法也是一样，所以对内存的开销很大，
 * 时间开销也大，所以这里只是提供一种实现的可能性，
 * 读者千万不要将其应用到实际中。如要应用，应先限定能要压缩的文件大小，如不超过20M等
 * @title: HuffmanTree
 * @author: zp
 * @date: 2021/12/14 19:59
 */
public class HuffmanTree {
    private static Map<Byte, String> huffmanMap = new HashMap<>();
    private static Map<String, Byte> codeMap = new HashMap<>();

    public static void main(String[] args) throws IOException, ClassNotFoundException {
        //压缩
        zipFile("E:\\Users\\aa.bmp", "E:\\Users\\aa.zip");
        //解压
        decodeZip("E:\\Users\\aa.zip", "E:\\Users\\cc.bmp");
    }

    /**
     * 解压文件
     *
     * @param zipPath
     * @param newPath
     */
    public static void decodeZip(String zipPath, String newPath) throws IOException, ClassNotFoundException {
        InputStream in = new FileInputStream(zipPath);
        ObjectInputStream ois = new ObjectInputStream(in);
        //读取byte数组
        byte[] filedatas = (byte[]) ois.readObject();
        //读取哈夫曼编码表
        Map<String, Byte> codeMap = (Map<String, Byte>) ois.readObject();
        //解码
        byte[] source = decodeByHuffmen(filedatas, codeMap);
        //byte[]输出到文件
        OutputStream out = new FileOutputStream(newPath);
        out.write(source);
        out.close();
    }

    /**
     * 使用哈夫曼编码进行文件压缩
     *
     * @param src 原文件地址
     * @param dst 压缩后的文件地址
     * @throws IOException
     */
    public static void zipFile(String src, String dst) throws IOException {
        InputStream in = new FileInputStream(src);
        byte[] srcBytes = new byte[in.available()];
        in.read(srcBytes);
        in.close();
        System.out.println("文件压缩前的大小：" + srcBytes.length);
        //使用哈夫曼压缩
        byte[] tarBytes = compressHuffman(srcBytes);
        System.out.println("文件压缩后的大小：" + tarBytes.length);

        //输出文件不仅包含压缩后的字节数据，还包含产生的哈夫曼编码，故使用包装类ObjectOutputStream
        OutputStream out = new FileOutputStream(dst);
        ObjectOutputStream oos = new ObjectOutputStream(out);
        oos.writeObject(tarBytes);
        oos.writeObject(codeMap);
        oos.close();
        out.close();
    }

    /**
     * @description: 压缩成哈夫曼字节流
     * @author: zp
     * @date: 2021/12/21 21:10
     * @param: bs
     * @return: byte[]
     */
    private static byte[] compressHuffman(byte[] bs) {
        if (bs == null || bs.length == 0) {
            return null;
        }
        //获取字符出现的次数和字符
        Map<Byte, Integer> countMap = getCharCount(bs);
        //构建哈夫曼树
        HuffmanNode huffmanNode = huffmanCoding(countMap);
        getHuCodeMap(huffmanNode, new StringBuilder());
        //返回转化哈夫曼编码组成的字符串
        List<String> strList = getHuffmanCodeList(bs);
        //转化压缩成byte数组
        return getByte(strList);
    }

    /**
     * 使用赫夫曼编码解压
     *
     * @param tar         目标数据
     * @param huffmenCode 赫夫曼编码表
     * @return
     */
    private static byte[] decodeByHuffmen(byte[] tar, Map<String, Byte> huffmenCode) {
        StringBuilder sb = new StringBuilder();
        //将byte[]转为二进制字符串
        for (int i = 0; i < tar.length; i++) {
            sb.append(byteToBitStr(tar[i]));
        }

        //根据哈夫曼编码表将二进制字符串转换成原数据
        List<Byte> source = new ArrayList<>();

        String strTotalCode = sb.toString();
        // 获取字符串总长度
        int strLength = strTotalCode.length();
        // 截取出最后八个之外的
        String strFact1 = strTotalCode.substring(0, strLength - 8);
        // 获取最后八个，并且转成对应的字节
        String lastEight = strTotalCode.substring(strLength - 8);
        // 得到补0的位数
        byte zeroNumber = bit2byte(lastEight);
        // 将得到的fact1减去最后的0的位数
        String strFact2 = strFact1.substring(0, strFact1.length() - zeroNumber);
        // 循环字节哈夫曼映射中的每一个哈夫曼值，然后在所有01串种进行匹配
        Set<Byte> byteSet = huffmanMap.keySet();
        int index = 0;
        // 从第0位开始
        String chs = strFact2.charAt(0) + "";
        while (index < strFact2.length()) {
            // 计数器，用来判断是否匹配到了
            int count = 0;
            for (Byte bi : byteSet) {
                // 如果匹配到了，则跳出循环
                if (chs.equals(huffmanMap.get(bi))) {
                    source.add(bi);
                    break;
                }
                // 没有匹配到则计数器累加一次
                count++;
            }
            // 如果计数器值大于或鱼等map，说明没有匹配到
            if (count >= byteSet.size()) {
                index++;
                chs += strFact2.charAt(index);
            } else {
                // 匹配到了，则匹配下一个字符串
                if (++index < strFact2.length()) {
                    chs = strFact2.charAt(index) + "";
                }
            }
        }

        //把集合转变成数组
        byte[] byteSource = new byte[source.size()];
        for (int i = 0; i < byteSource.length; i++) {
            byteSource[i] = source.get(i);
        }
        return byteSource;
    }

    /**
     * 将字符串转成二进制字节的方法
     *
     * @param bString 待转换的字符串
     * @return 二进制字节
     */
    private static byte bit2byte(String bString) {
        byte result = 0;
        for (int i = bString.length() - 1, j = 0; i >= 0; i--, j++) {
            result += (Byte.parseByte(bString.charAt(i) + "") * Math.pow(2, j));
        }
        return result;
    }

    /**
     * @description: 把字节转化回字符串 （补码的形式 补码=源码的反码+1）
     * @author: zp
     * @date: 2021/12/21 21:21
     * @param: b
     * @return: java.lang.String
     */
    private static String byteToBitStr(byte b) {
        //将8位扩大到32位，便于 或 运算，提取原来数值中的8位
        int temp = b;
        temp |= 256;
        String str = Integer.toBinaryString(temp);
        return str.substring(str.length() - 8);
    }

    /**
     * @description: 转化之后的数据压缩
     * @author: zp
     * @date: 2021/12/20 21:00
     * @param: strList
     * @return: byte[]
     */
    public static byte[] getByte(List<String> strList) {
        StringBuilder stringBuilder = new StringBuilder();
        for (String str : strList) {
            stringBuilder.append(str);
        }
        String s = stringBuilder.toString();
        int zeroNumber = 0;
        int len = (s.length() + 7) / 8;
        byte[] bytes = new byte[len + 1];
        int index = 0;
        for (int i = 0; i < s.length(); i += 8) {
            String strByte;
            if (i + 8 > s.length()) {
                strByte = s.substring(i);
                zeroNumber = 8 - strByte.length();
                for (int c = 0; c < zeroNumber; c++) {
                    // 补0操作
                    strByte += "0";
                }

            } else {
                strByte = s.substring(i, i + 8);
            }
            bytes[index] = (byte) Integer.parseInt(strByte, 2);
            index++;
        }
        String zeroLenStr = Integer.toBinaryString(zeroNumber);
        bytes[bytes.length - 1] = (byte) Integer.parseInt(zeroLenStr, 2);
        return bytes;
    }

    /**
     * @description: 字符串转换成赫弗曼码
     * @author: zp
     * @date: 2021/12/18 16:09
     * @param: s
     * @return: java.util.List<java.lang.String>
     */
    public static String getHuffmanCodeStr(List<String> list) {
        if (list == null || list.size() == 0) {
            return null;
        }
        StringBuilder stringBuilder = new StringBuilder();

        for (String s : list) {
            stringBuilder.append(codeMap.get(s));
        }

        return stringBuilder.toString();
    }

    /**
     * @description: 字节数组转换成赫弗曼码
     * @author: zp
     * @date: 2021/12/18 16:09
     * @param: s
     * @return: java.util.List<java.lang.String>
     */
    public static List<String> getHuffmanCodeList(byte[] bs) {
        if (bs == null || bs.length == 0) {
            return null;
        }
        List<String> list = new ArrayList<>();

        for (byte c : bs) {
            list.add(huffmanMap.get(c));
        }

        return list;
    }


    /**
     * @description: 求哈夫曼编码的算法, W存放n个字符的权值（均 > 0）
     * @author: zp
     * @date: 2021/12/18 15:48
     * @param: map
     * @return: tree.HuffmanNode
     */
    public static HuffmanNode huffmanCoding(Map<Byte, Integer> map) {
        //定义一个霍夫曼节点的数组
        List<HuffmanNode> nodeList = new ArrayList<>();

        map.forEach((key, value) -> {
            //构造n个具有权值的节点，因为并不是每一个节点都是有权值的，只有在边上的才有
            nodeList.add(new HuffmanNode(value, key));
        });

        //开始建立哈夫曼树
        while (nodeList.size() > 1) {
            Collections.sort(nodeList);
            //在 nodeList[0..i-1] 选择不在哈夫曼树中且weight最小的两个节点 min1 和 min2
            HuffmanNode min1 = nodeList.get(0);
            HuffmanNode min2 = nodeList.get(1);


            //构造min1和min2的父节点，并修改其父节点的权值
            HuffmanNode huffmanNode = new HuffmanNode();
            huffmanNode.lChild = min1;
            huffmanNode.rChild = min2;
            huffmanNode.flag = 0;
            huffmanNode.weight = min1.weight + min2.weight;
            nodeList.remove(0);
            nodeList.remove(0);
            nodeList.add(huffmanNode);
        }

        return nodeList.get(0);
    }

    /**
     * @description: 转化为赫弗曼编码表
     * @author: zp
     * @date: 2021/12/18 15:56
     * @param: node
     * @param: stringBuilder
     * @return: void
     */
    public static void getHuCodeMap(HuffmanNode node, StringBuilder stringBuilder) {
        if (node == null) {
            return;
        }
        StringBuilder leftStr = new StringBuilder();
        getHuCodeMap(node.lChild, leftStr.append(stringBuilder).append("0"));
        if (node.flag == 1) {
            huffmanMap.put(node.data, stringBuilder.toString());
            codeMap.put(stringBuilder.toString(), node.data);
        }
        StringBuilder rightStr = new StringBuilder();
        getHuCodeMap(node.rChild, rightStr.append(stringBuilder).append("1"));
    }

    /**
     * 将bytes的中的元素以及出现次数包装成HuffmanNode列表
     *
     * @param bytes
     * @return
     */
    private static Map<Byte, Integer> getCharCount(byte[] bytes) {
        Map<Byte, Integer> byteIntegerMap = new HashMap<>();
        for (byte b : bytes) {
            Integer count = byteIntegerMap.get(b);
            if (count == null) {
                byteIntegerMap.put(b, 1);
            } else {
                byteIntegerMap.put(b, count + 1);
            }
        }
        return byteIntegerMap;
    }
}

